from collections.abc import Iterable, Mapping, Sequence
from typing import TypeVar

from zkl_aiutils_datasets.basics import Dataset, GroupedDatasets, IndexedDatasets, NamedDatasets
from zkl_aiutils_datasets.utils.lazy import LazySequence

AnySample = TypeVar('AnySample')
AnyDataset = TypeVar('AnyDataset', bound=Dataset)


class SimpleGroupedDatasets(GroupedDatasets[AnyDataset]):
    def __init__(self, children: Iterable[AnyDataset]):
        self._children = children

    @property
    def children(self) -> Iterable[AnyDataset]:
        return self._children


class SimpleIndexedDatasets(IndexedDatasets[AnyDataset]):
    def __init__(self, indexed_children: Sequence[AnyDataset]):
        self._indexed_children = indexed_children

    @property
    def indexed_children(self) -> Sequence[AnyDataset]:
        return self._indexed_children


class SimpleNamedDatasets(NamedDatasets[AnyDataset]):
    def __init__(self, named_children: Mapping[str, AnyDataset]):
        self._named_children = named_children

    @property
    def named_children(self) -> Mapping[str, AnyDataset]:
        return self._named_children


class SimpleIndexedNamedDatasets(SimpleIndexedDatasets[AnyDataset], SimpleNamedDatasets[AnyDataset]):
    def __init__(self, children: Mapping[str, Dataset[AnySample]] | Iterable[tuple[str, Dataset[AnySample]]]):
        if isinstance(children, Mapping):
            child_getter_factory = lambda name: lambda: children[name]
            children_getter = tuple(child_getter_factory(name) for name in children.keys())
            SimpleIndexedDatasets.__init__(self, LazySequence(children_getter))
            SimpleNamedDatasets.__init__(self, children)
        else:
            SimpleIndexedDatasets.__init__(self, tuple(child for _, child in children))
            SimpleNamedDatasets.__init__(self, dict(children))
